package com.kakalog.main.gpt;

import java.net.Proxy;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

import com.alibaba.fastjson2.JSON;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;


@Slf4j
public class GptRequestClient {

    @Getter
    @NotNull
    private String apiKey;

    /**
     * 自定义api host使用builder的方式构造client
     */
    @Getter
    private String apiHost = "https://usechat.cn/";

    @Getter
    private OkHttpClient okHttpClient;

    /**
     * 连接超时
     */
    @Getter
    private long connectTimeout;

    /**
     * 写超时
     */
    @Getter
    private long writeTimeout;

    /**
     * 读超时
     */
    @Getter
    private long readTimeout;

    /**
     * okhttp 代理
     */
    @Getter
    private Proxy proxy;

    public GptRequestClient(@NotNull String apiKey) {
        this.apiKey = apiKey;
        this.okHttpClient(30, 30, 30, null);
    }

    public GptRequestClient(@NotNull String apiKey, Proxy proxy) {
        this.apiKey = apiKey;
        this.proxy = proxy;
        this.okHttpClient(30, 30, 30, proxy);
    }

    public GptRequestClient(@NotNull String apiKey, long connectTimeout, long writeTimeout, long readTimeout, Proxy proxy) {
        this.apiKey = apiKey;
        this.connectTimeout = connectTimeout;
        this.writeTimeout = writeTimeout;
        this.readTimeout = readTimeout;
        this.proxy = proxy;
        okHttpClient(connectTimeout, writeTimeout, readTimeout, proxy);
    }

    /**
     * 创建 OkHttpClient，自定义超时时间和拦截器
     * @param connectTimeout 默认30秒
     * @param writeTimeout   默认30秒
     * @param readTimeout    默认30秒
     */
    private void okHttpClient(long connectTimeout, long writeTimeout, long readTimeout, Proxy proxy) {
        OkHttpClient.Builder client = new OkHttpClient.Builder();
        client.connectTimeout(connectTimeout, TimeUnit.SECONDS);
        client.writeTimeout(writeTimeout, TimeUnit.SECONDS);
        client.readTimeout(readTimeout, TimeUnit.SECONDS);
        if (Objects.nonNull(proxy)) {
            client.proxy(proxy);
        }
        this.okHttpClient = client.build();
    }

    private GptRequestClient(Builder builder) {
        if (StringUtils.isBlank(builder.apiKey)) {
            throw new IllegalArgumentException();
        }
        apiKey = builder.apiKey;

        if (StringUtils.isBlank(builder.apiHost)) {
            builder.apiHost = "https://api.openai.com/";
        }
        apiHost = builder.apiHost;

        if (Objects.isNull(builder.connectTimeout)) {
            builder.connectTimeout(30);
        }
        connectTimeout = builder.connectTimeout;

        if (Objects.isNull(builder.writeTimeout)) {
            builder.writeTimeout(30);
        }
        writeTimeout = builder.writeTimeout;

        if (Objects.isNull(builder.readTimeout)) {
            builder.readTimeout(30);
        }
        readTimeout = builder.readTimeout;
        proxy = builder.proxy;
        this.okHttpClient(connectTimeout, writeTimeout, readTimeout, proxy);
    }

    public static final class Builder {
        private @NotNull String apiKey;
        private String apiHost;
        private long connectTimeout;
        private long writeTimeout;
        private long readTimeout;
        private Proxy proxy;

        public Builder() {
        }

        public Builder apiKey(@NotNull String val) {
            apiKey = val;
            return this;
        }

        public Builder apiHost(String val) {
            apiHost = val;
            return this;
        }

        public Builder connectTimeout(long val) {
            connectTimeout = val;
            return this;
        }

        public Builder writeTimeout(long val) {
            writeTimeout = val;
            return this;
        }

        public Builder readTimeout(long val) {
            readTimeout = val;
            return this;
        }

        public Builder proxy(Proxy val) {
            proxy = val;
            return this;
        }

        public GptRequestClient build() {
            return new GptRequestClient(this);
        }
    }

    /**
     * 构造
     */
    public static Builder builder() {
        return new Builder();
    }

    public void streamChatCompletion(ChatGPTParam chatGPTParam, EventSourceListener eventSourceListener) {
        if (Objects.isNull(eventSourceListener)) {
            throw new IllegalArgumentException();
        }
        try {
            EventSource.Factory factory = EventSources.createFactory(this.okHttpClient);
            String requestBody = JSON.toJSONString(chatGPTParam);
            LogUtil.info("request gpt requestBody: ", requestBody);
            Request request = new Request.Builder()
                .url(this.apiHost + "v1/chat/completions")
                .post(RequestBody.create(MediaType.parse("application/json"), requestBody))
                .header("Authorization", "Bearer " + this.apiKey)
                .build();
            //创建事件
            factory.newEventSource(request, eventSourceListener);
        } catch (Exception e) {
            LogUtil.error("GptRequestClient#streamChatCompletion error.", e);
        }
    }

}
